import json
from argparse import ArgumentParser
import torch as th
import torch
from torch import nn
from torchvision.models.resnet import resnet18, resnet34, ResNet18_Weights, ResNet34_Weights
from skimage import io
from preprocess import preprocess

print('torch version:', th.__version__)
print('cuda:', th.cuda.is_available())
torch.backends.cudnn.benchmark = True


with open('class_id.json') as f:
    class_id = json.load(f)

parser = ArgumentParser(description='run tensorrt resnet on image')
parser.add_argument('images', type=str, nargs='+')
ARGS = parser.parse_args()
images = ARGS.images

network = resnet18(weights=None)
network.eval()
network.cuda()

'''
with th.no_grad():
    #layer = nn.Linear(10, 10)
    layer = nn.Conv2d(3, 1, kernel_size=3, stride=1, padding=1, bias=False)
    test_input = th.randn((1, 3, 224, 224), dtype=th.float32)
    # test_input = th.randn((224, 10), dtype=th.float32)
    test_output = layer(test_input)
    print('test_output.max:', th.max(test_output))
    print('test_output.min:', th.min(test_output))
'''


for filename in images:
    image = io.imread(filename)
    image = preprocess(image)
    print('process image:', filename, 'shape:', image.shape)
    input = th.tensor(image)
    input.unsqueeze_(0)
    assert th.isnan(input).sum() == 0
    print('input.shape:', input.shape)

    with th.no_grad():
        input = th.ones((1, 3, 224, 224), dtype=th.float32)
        input = input.cuda()
        # assert th.isnan(layer1).sum() == 0
        output = network(input)
        print('output.shape:', output.shape)
        print('output[0, :10]:', output[0, :10])

    assert th.isnan(output).sum() == 0

    infer_index = th.argmax(output)
    print('infer_index:', infer_index)
    print('infer class:', class_id[infer_index])

